import cv2
import os
import numpy as np

# 初始化人脸检测器和识别器
face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
recognizer = cv2.face.LBPHFaceRecognizer_create()

def get_dataset(path):
    faces = []
    labels = []
    label_dict = {}
    current_label = 0

    # 遍历数据集目录
    for root, dirs, files in os.walk(path):
        for dir_name in dirs:
            label_dict[current_label] = dir_name
            subject_path = os.path.join(root, dir_name)

            # 处理每个人的图片
            for file in os.listdir(subject_path):
                img_path = os.path.join(subject_path, file)
                img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)

                # 检测人脸并裁剪
                faces_rect = face_cascade.detectMultiScale(img)
                for (x, y, w, h) in faces_rect:
                    face = img[y:y+h, x:x+w]
                    faces.append(face)
                    labels.append(current_label)

            current_label += 1

    return faces, labels, label_dict

# 数据集路径
dataset_path = 'dataset'

# 获取数据集和标签
faces, labels, label_dict = get_dataset(dataset_path)

# 训练模型
recognizer.train(faces, np.array(labels))

# 保存模型和标签
recognizer.save('face_model.yml')
np.save('label_dict.npy', label_dict)